import os
from glob import glob

import cv2
import numpy as np
from tqdm import tqdm


def main():
    """
    预处理ISIC 2018数据集
    将原始数据转换为统一格式：
    inputs/isic2018_{size}/
        ├── images/
        └── masks/0/
    """
    
    # 可选择的图像尺寸
    img_size = 256  # 推荐256x256，ISIC图像原始尺寸较大
    
    # 输入路径
    train_img_dir = 'inputs/ISIC2018_Task1-2_Training_Input'
    train_mask_dir = 'inputs/ISIC2018_Task1_Training_GroundTruth'
    val_img_dir = 'inputs/ISIC2018_Task1-2_Validation_Input'
    val_mask_dir = 'inputs/ISIC2018_Task1_Validation_GroundTruth'
    
    # 输出路径
    output_dir = f'inputs/isic2018_{img_size}'
    os.makedirs(f'{output_dir}/images', exist_ok=True)
    os.makedirs(f'{output_dir}/masks/0', exist_ok=True)
    
    print(f'开始处理ISIC 2018数据集，目标尺寸: {img_size}x{img_size}')
    
    # 处理训练集
    print('\n处理训练集...')
    train_imgs = glob(os.path.join(train_img_dir, '*.jpg'))
    
    processed_count = 0
    skipped_count = 0
    
    for img_path in tqdm(train_imgs, desc='训练集'):
        # 获取文件名（去掉扩展名）
        basename = os.path.basename(img_path).replace('.jpg', '')
        
        # 对应的mask文件名
        mask_path = os.path.join(train_mask_dir, f'{basename}_segmentation.png')
        
        # 检查mask是否存在
        if not os.path.exists(mask_path):
            skipped_count += 1
            continue
        
        # 读取图像和mask
        img = cv2.imread(img_path)
        mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)
        
        if img is None or mask is None:
            skipped_count += 1
            continue
        
        # 确保图像是RGB格式
        if len(img.shape) == 2:
            img = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)
        elif img.shape[2] == 4:
            img = cv2.cvtColor(img, cv2.COLOR_BGRA2BGR)
        
        # 调整大小
        img = cv2.resize(img, (img_size, img_size), interpolation=cv2.INTER_LINEAR)
        mask = cv2.resize(mask, (img_size, img_size), interpolation=cv2.INTER_NEAREST)
        
        # 二值化mask（确保只有0和255）
        mask = (mask > 127).astype(np.uint8) * 255
        
        # 保存
        cv2.imwrite(f'{output_dir}/images/{basename}.png', img)
        cv2.imwrite(f'{output_dir}/masks/0/{basename}.png', mask)
        
        processed_count += 1
    
    print(f'训练集: 已处理 {processed_count} 张, 跳过 {skipped_count} 张')
    
    # 处理验证集
    print('\n处理验证集...')
    val_imgs = glob(os.path.join(val_img_dir, '*.jpg'))
    
    val_processed = 0
    val_skipped = 0
    
    for img_path in tqdm(val_imgs, desc='验证集'):
        basename = os.path.basename(img_path).replace('.jpg', '')
        mask_path = os.path.join(val_mask_dir, f'{basename}_segmentation.png')
        
        if not os.path.exists(mask_path):
            val_skipped += 1
            continue
        
        img = cv2.imread(img_path)
        mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)
        
        if img is None or mask is None:
            val_skipped += 1
            continue
        
        if len(img.shape) == 2:
            img = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)
        elif img.shape[2] == 4:
            img = cv2.cvtColor(img, cv2.COLOR_BGRA2BGR)
        
        img = cv2.resize(img, (img_size, img_size), interpolation=cv2.INTER_LINEAR)
        mask = cv2.resize(mask, (img_size, img_size), interpolation=cv2.INTER_NEAREST)
        mask = (mask > 127).astype(np.uint8) * 255
        
        cv2.imwrite(f'{output_dir}/images/{basename}.png', img)
        cv2.imwrite(f'{output_dir}/masks/0/{basename}.png', mask)
        
        val_processed += 1
    
    print(f'验证集: 已处理 {val_processed} 张, 跳过 {val_skipped} 张')
    
    print(f'\n✓ 预处理完成！')
    print(f'总计: {processed_count + val_processed} 张图像')
    print(f'输出目录: {output_dir}')
    print(f'\n下一步运行训练命令:')
    print(f'python train.py --dataset isic2018_{img_size} --arch EnhancedNestedUNet --img_ext .png --mask_ext .png')


if __name__ == '__main__':
    main()






